



# -*- coding: utf-8 -*-
import os
import base64
import logging
import tempfile
from configparser import ConfigParser
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes
from typing import Optional, Tuple

from PyQt5.QtWidgets import QHBoxLayout, QPushButton, QVBoxLayout, QLabel, QLineEdit, QWidget, QScrollArea, QDialog, \
    QMessageBox
from cryptography.fernet import Fernet

# 配置日志
logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger(__name__)


class EncryptionError(Exception):
    """加密操作异常"""
    pass


class DecryptionError(Exception):
    """解密操作异常"""
    pass


class ConfigEncryptor:
    def __init__(self, config_path: str = 'config.ini',
                 key_path: Optional[str] = None):
        """
        :param config_path: 配置文件路径
        :param key_path: 密钥文件路径（默认用户目录）
        """
        self.config_path = config_path
        self.key_path = key_path or os.path.expanduser('~/.app_secure.key')
        self.config = ConfigParser()
        self._init_config()
        self._load_or_generate_key()
        self._validate_keys()

    def _init_config(self):
        """初始化配置文件"""
        """初始化配置文件（显式指定UTF-8编码）"""
        if not os.path.exists(self.config_path):
            self._create_default_config()
        # 修复：添加 encoding='utf-8' 参数
        with open(self.config_path, 'r', encoding='utf-8') as f:
            self.config.read_file(f)  # 使用 read_file 替代 read
        if not self.config.has_section('Security'):
            self.config.add_section('Security')

    def _create_default_config(self):
        """创建默认配置文件（使用 UTF-8 编码）"""
        with open(self.config_path, 'w', encoding='utf-8') as f:
            self.config['DEFAULT'] = {
                # 现在 Fernet 可以被正确识别
                'encryption_key': base64.b64encode(get_random_bytes(32)).decode('utf-8'),  # 生成AES-256密钥
                'cookie_lifetime': '30'  # 默认30天
            }
            self.config.write(f)

    def _load_or_generate_key(self):
        """加载或生成加密密钥"""
        if os.path.exists(self.key_path):
            self._load_key()
        else:
            self._generate_key()
            self._save_key()

    def _validate_keys(self):
        """验证密钥有效性"""
        if len(self.key) != 32:
            logger.error("密钥长度无效，应为32字节（AES-256）")
            self._generate_key()  # 重新生成密钥
            self._save_key()
            logger.warning("已自动生成新密钥")

    def _generate_key(self):
        """生成符合 Fernet 要求的密钥（32 bytes base64编码）"""
        self.key = Fernet.generate_key()  # 直接使用 Fernet 生成密钥
        logger.info("生成新的加密密钥")

    def _save_key(self):
        """保存密钥文件"""
        key_dir = os.path.dirname(self.key_path)
        os.makedirs(key_dir, exist_ok=True)

        with open(self.key_path, 'wb') as f:
            f.write(self.key)
        os.chmod(self.key_path, 0o600)  # 设置文件权限
        logger.info(f"密钥已保存至 {self.key_path}")

    def _load_key(self):
        """加载现有密钥"""
        with open(self.key_path, 'rb') as f:
            self.key = f.read()
        logger.info("加密密钥已加载")

    def encrypt_value(self, plaintext: str) -> str:
        if not isinstance(plaintext, str):
            raise ValueError("输入必须是字符串")

        try:
            # 限制输入大小（防止超大内存分配）
            if len(plaintext) > 1024 * 1024:  # 1MB限制
                raise ValueError("Cookie值过长")

            iv = get_random_bytes(16)
            cipher = AES.new(self.key, AES.MODE_CBC, iv)
            padded = pad(plaintext.encode('utf-8'), AES.block_size)
            ciphertext = cipher.encrypt(padded)
            return base64.b64encode(iv + ciphertext).decode('ascii')  # 强制ASCII输出
        except Exception as e:
            logger.error(f"加密失败: {type(e).__name__}: {str(e)}")
            raise EncryptionError(f"加密错误: {str(e)}") from e

    def decrypt_value(self, ciphertext_b64: str) -> str:
        """解密加密字符串"""
        try:
            data = base64.b64decode(ciphertext_b64)
            iv, ciphertext = data[:16], data[16:]
            cipher = AES.new(self.key, AES.MODE_CBC, iv)
            plaintext_padded = cipher.decrypt(ciphertext)
            return unpad(plaintext_padded, AES.block_size).decode('utf-8')
        except Exception as e:
            logger.error(f"解密失败: {str(e)}")
            raise DecryptionError("值解密失败") from e

    def secure_save(self, section: str, key: str, value: str):
        """安全保存配置项"""
        encrypted = self.encrypt_value(value)
        self.config.set(section, key, encrypted)
        self.save_config()

    def secure_get(self, section: str, key: str,
                   default: Optional[str] = None) -> Optional[str]:
        """安全读取配置项"""
        try:
            encrypted = self.config.get(section, key)
            return self.decrypt_value(encrypted)
        except (DecryptionError, Exception):
            return default

    def save_config(config):
        """
        线程安全的配置文件保存方法
        :param config: ConfigParser 对象
        :return: bool 是否成功
        """
        temp_path = None
        try:
            # 创建临时文件
            with tempfile.NamedTemporaryFile(
                    mode='w',
                    encoding='utf-8',
                    delete=False,  # 不自动删除
                    suffix='.tmp',
                    prefix='cookies_'
            ) as tmp:
                temp_path = tmp.name
                config.write(tmp)

            # 原子性替换原文件
            if os.path.exists('cookies.ini'):
                os.replace(temp_path, 'cookies.ini')
            else:
                os.rename(temp_path, 'cookies.ini')

            return True

        except Exception as e:
            # 清理临时文件
            if temp_path and os.path.exists(temp_path):
                try:
                    os.unlink(temp_path)
                except:
                    pass
            logger.error(f"保存配置失败: {str(e)}")
            return False
        finally:
            # 确保文件句柄关闭
            if 'tmp' in locals():
                tmp.close()

    def rotate_keys(self):
        """轮换加密密钥并更新所有加密值"""
        old_key = self.key
        self._generate_key()

        # 重新加密所有配置项
        for section in self.config.sections():
            for key in self.config[section]:
                try:
                    value = self.decrypt_value(self.config[section][key])
                    self.config[section][key] = self.encrypt_value(value)
                except DecryptionError:
                    continue

        self._save_key()
        logger.warning("加密密钥已轮换，请妥善保管新密钥")




# 示例用法
if __name__ == '__main__':
    # 初始化加密器
    encryptor = ConfigEncryptor()

    # 保存敏感配置
    encryptor.secure_save('Credentials', 'api_key', 'my_secret_key_123')

    # 读取加密配置
    api_key = encryptor.secure_get('Credentials', 'api_key')
    print(f"解密后的API密钥: {api_key}")

    # 测试密钥轮换
    # encryptor.rotate_keys()



